import socket
import json
import threading
import time
import cv2
import numpy as np
import math
import queue
class GlimmerNode:
    father=None
    child_n=None#是父节点的第几个子节点(从1开始)
    n_child=0#有几个子节点
    depth=0#该节点的深度(在深度优先遍历的时候更新,如果树的结构变化，该值可能不对，需要更新)
    child=[]
    child_name=[]
    name=""
    #----------------画图相关------------
    x=None
    y=None
    state=0
    def __init__(self,name,father):
        self.name=name
        self.father=father
        self.child=[]       #有多个实例时不初始化参数会出错
        self.child_name=[]
        self.n_child=0
        self.depth=0
        if father:
            self.depth=father.depth+1
        return
    def __str__(self):
        father=self.father
        if father==None:
            father="None"
        else:
            father=father.name
        return "name:{:10s} n_child:{} child_name:{} father:{:10s} child_n:{}".format(self.name,self.n_child,self.child_name,father,self.child_n)
class GlimmerTree:
    #--------------基本数据--------------
    root=None
    depth_all=0
    child_n_max=0
    names=[]#节点的名字,按添加的顺序
    nodes=[]#节点实例，按添加的顺序
    deep_first_names=[]#深度优先遍历后的名字
    deep_first_nodes=[]#深度优先遍历后的节点实例
    #--------------画图相关--------------
    depth=70  #节点方块高
    width=270 #节点方块长
    delta_x=300#节点方块横向距离
    delta_y=150#节点方块纵向距离
    canvas=None
    th_show=None
    running=True
    window_index=1
    #--------------基本操作--------------
    def __init__(self):
        self.root=None
        self.names=[]#有多个实例时不初始化参数会出错
        self.nodes=[]
        self.deep_first_names=[]#深度优先遍历后的名字
        self.deep_first_nodes=[]#深度优先遍历后的节点实例

        return
    def add_root(self,name):
        if self.root!=None:
            raise RuntimeError("This tree already have a root node")
        self.root=GlimmerNode(name,None)
        self.names.append(name)
        self.nodes.append(self.root)
        self.deep_first_names.append(name)
        self.deep_first_nodes.append(self.root)
        return self.root
    def clear_buffer(self):
        self.deep_first_names=[]
        self.deep_first_nodes=[]
        self.depth_all=0
        self.child_n_max=0
        self.deep_first_names.append(self.root.name)
        self.deep_first_nodes.append(self.root)
        return
    def deep_first_traversal(self,Node):#深度优先遍历,遍历前需要clear_buffer
        if Node==None:
            raise ValueError("deep_first_traversal received a None node")
            return
        for node in Node.child:
            node.depth=Node.depth+1
            self.deep_first_names.append(node.name)
            self.deep_first_nodes.append(node)
            self.deep_first_traversal(node)
        return
    def initialize_parameters(self):#在深度遍历之后，获取树的深度等参数
        current_depth=0
        for i in range(0,len(self.deep_first_nodes)):
            node=self.deep_first_nodes[i]
            if node.n_child>self.child_n_max:
                self.child_n_max=node.n_child
            if node.depth>self.depth_all:
                self.depth_all=node.depth
        return
    def add_node(self,name,father):
        new_node=GlimmerNode(name,father)
        father.n_child+=1
        father.child.append(new_node)
        father.child_name.append(name)
        new_node.child_n=father.n_child
        self.names.append(name)
        self.nodes.append(new_node)
        return new_node
    def find(self,name):#只能查找到用GlimmerTree添加的节点
        if name in self.names:
            index=self.names.index(name)
        else:
            return -1
        return self.nodes[index]
    def deep_first_find(self,name):
        self.clear_buffer()
        self.deep_first_traversal(self.root)

        if name in self.deep_first_names:
            index=self.deep_first_names.index(name)
        else:
            return -1
        return self.deep_first_nodes[index]
    #--------------画图相关--------------
    def creat_window(self):
        cv2.destroyAllWindows()
        window_name="topology"+str(self.window_index)
        cv2.namedWindow(window_name,cv2.WINDOW_AUTOSIZE) #创建窗口
        cv2.moveWindow(window_name,100,100)
        return
    def refresh(self):
        self.clear_buffer()
        self.deep_first_traversal(self.root)
        self.initialize_parameters()
        canvas_x=math.ceil(self.delta_y*self.depth_all+self.depth)
        canvas_y=math.ceil(self.delta_x*(math.pow(self.child_n_max,self.depth_all)-1)+self.width)
        canvas_shape=(canvas_x,canvas_y,3)
        canvas = np.zeros(canvas_shape, dtype="uint8")
        self.root.x=self.delta_x*(math.pow(self.child_n_max,self.depth_all)-1)/2
        self.root.y=0
        for i in range(1,len(self.deep_first_nodes)):
            node=self.deep_first_nodes[i]
            delta_x=self.delta_x*math.pow(2,self.depth_all-node.depth)
            node.x=node.father.x+delta_x*(node.child_n-0.5-self.child_n_max/2)
            node.y=node.father.y+self.delta_y
        for i in range(0,len(self.deep_first_nodes)):
            node=self.deep_first_nodes[i]
            cv2.rectangle(canvas,(round(node.x),round(node.y)),(round(node.x+self.width),round((node.y+self.depth))),(255,255,255),1)  #画矩形
            if i>0:
                cv2.line(canvas, (round(node.x+self.width/2),round(node.y)),(round(node.father.x+self.width/2),round(node.father.y+self.depth)),(255, 255, 255),1 )
            text_depth=30
            cv2.putText(canvas,node.name,(round(node.x),round(node.y+text_depth)),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,255),2)
            cv2.putText(canvas,"state:"+str(node.state),(round(node.x),round(node.y+2*text_depth)),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,255),2)
        self.canvas=canvas
        return
    def opencv_show(self):
        self.creat_window()
        canvas=self.refresh()
        window_name="topology"+str(self.window_index)
        while self.running:
            cv2.imshow(window_name, self.canvas) #10
            cv2.waitKey(50)
        return
    def start_show_topology(self,window_index):
        self.window_index=window_index
        self.th_show= threading.Thread(target=self.opencv_show, args=())
        self.th_show.start()
        return
    def stop(self):
        self.running=False
        self.th_show.join()
        #cv2.destroyAllWindows()
        if self.root:
            self.clear_buffer()
            self.deep_first_traversal(self.root)
        for node in self.deep_first_nodes:
            del node
        return
class FLcloud:
    #----------基本控制-------------
    waite_time=3.5#超过这个时间没有收到信息就认为是退出了
    command_id=set()
    port=23334
    local_ip="192.168.118.1"
    FL_server_n=0
    FL_server_ip=[]#["192.168.118.130","192.168.0.2","192.168.0.3","192.168.0.4"]
    FL_server_isroot=[]#[False,True,False,False]
    FL_server_state=[]#[0,0,0,0]
    FL_server_up=[]#["192.168.0.2","","192.168.0.2","192.168.0.2"]#FL节点的上层节点
    FL_server_down=[]#[[],["192.168.118.130","192.168.0.3","192.168.0.4"],[],[]]#FL节点的下层节点
    FL_server_topology_changed=[]
    FL_server_timmer=[]
    tree=None#拓扑关系
    flag_initialized=False#第一次获取完整拓扑结构时变为真
    sendjson1=json.dumps({"order":"aggregation_request"})
    sendjson3=json.dumps({"order":"test_local_model"})
    sendjson4=json.dumps({"order":"save_model"})
    sendjson5=json.dumps({"order":"read_model"})
    sendjson6=json.dumps({"order":"update_part1"})
    
    order={"1":"start a round of FL",
           "2":"change topology",
           "3":"test local model",
           "4":"save local model",
           "5":"read local model",
           "6":"update",
           "7":"multipule turns auto test",
           "0":"refresh"}
    running=True
    window_index=0
    def __init__(self):
        return
    def listen(self):
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        server_socket.bind(("", self.port)) 
        while(self.running):
            server_socket.settimeout(3)  #每3秒钟检测是否关闭
            try:
                receive_data, client = server_socket.recvfrom(1400)
                #print("来自客户端%s,发送的%s\n" % (client, receive_data))  #打印接收的内容
                josn_end=receive_data.index(b"}")+1
                receive_data2=receive_data[josn_end:]
                json_data = json.loads(receive_data[:josn_end].decode('utf-8'))
                if "feedback_id" in json_data:
                    if not self.feedback(json_data):
                        json_data={}#已经处理过的任务就跳过
                if 'order' in json_data:
                    if json_data['order']=="send_state":
                        th_receive_state= threading.Thread(target=self.receive_state, args=(json_data,))
                        th_receive_state.start()
                    elif json_data['order']=="print":
                        print(json_data['text'])
                    elif json_data['order']=="feedback":
                        if "detail" in json_data:
                            if json_data['detail']=="topology_changed":
                                ip_index=self.FL_server_ip.index(json_data['ip'])
                                self.FL_server_topology_changed[ip_index]=True
            except socket.timeout:
                pass
        return
    def show_state(self):
        while self.running:
            print('\r',end='')
            for i in range(self.FL_server_n):
                print(" ip:{} state:{} isroot:{} up_server:{} down_server{}".format(self.FL_server_ip[i],self.FL_server_state[i],self.FL_server_isroot[i],self.FL_server_up[i],self.FL_server_down[i]),end='')
            time.sleep(0.5)
        return
    def time_task(self):
        def check_server_timmer():
            current_time=time.time()
            for i in range(self.FL_server_n):
                if current_time-self.FL_server_timmer[i]>self.waite_time:
                    print("FL server:{} time out".format(self.FL_server_ip[i]))
                    self.FL_server_n-=1
                    del self.FL_server_ip[i]
                    del self.FL_server_state[i]
                    del self.FL_server_isroot[i]
                    del self.FL_server_up[i]
                    del self.FL_server_down[i]
                    del self.FL_server_timmer[i]
                    del self.FL_server_topology_changed[i]
                    while not self.check_state():
                        print("\rwaite...",end='')
                        time.sleep(1)
                    if self.FL_server_n>0:
                        self.topology_test()
                    else:
                        self.tree.stop()
                        self.flag_initialized=False
                        print("No FL server")
                    break
            return
        while self.running:
            check_server_timmer()
            time.sleep(1)
        return
    def change_topology(self,ip,up_server,down_server):
        sendjson=json.dumps({"order":"change_topology","up_server":up_server,"down_server":down_server})
        return self.command_send((ip,self.port),sendjson)
    def start(self):
        for i in self.order:
            print(i,self.order[i])
        self.th_listen= threading.Thread(target=self.listen, args=())
        self.th_listen.start()
        #self.th_show= threading.Thread(target=self.show_state, args=())
        #self.th_show.start()
        self.th_time_task= threading.Thread(target=self.time_task, args=())
        self.th_time_task.start()
        th_inpute=threading.Thread(target=self.keyboard_input, args=())
        th_inpute.start()
        return
    def keyboard_input(self):
        while self.running:
            a = input()
            if a=="1":
                root_index=self.FL_server_isroot.index(True)
                self.command_send((self.FL_server_ip[root_index],self.port),self.sendjson1)
            elif a=="2":
                self.topology_test()
            elif a=="3":
                for ip in self.FL_server_ip:
                    self.command_send((ip,self.port),self.sendjson3)
            elif a=="0":
                self.refresh()
            elif a=="4":
                for ip in self.FL_server_ip:
                    self.command_send((ip,self.port),self.sendjson4)
            elif a=="5":
                for ip in self.FL_server_ip:
                    self.command_send((ip,self.port),self.sendjson5)
            elif a=="6":
                for ip in self.FL_server_ip:
                    self.command_send((ip,self.port),self.sendjson6)
            elif a=="7":
                self.multipule_test()
            elif a=="q":
                self.running=False
        return
    def refresh(self):#更新显示的拓扑信息
        self.window_index+=1
        if self.tree!=None:
            self.tree.stop()
            del self.tree
        self.tree=GlimmerTree()
        first_time=True
        def tree_grow(tree,index,first_time):#生成树，输入树和root节点的index
            if first_time:
                first_time=False
                Node=tree.add_root(self.FL_server_ip[index])
                tree_grow(tree,index,first_time)
            else:
                for server in self.FL_server_down[index]:
                    index_new=self.FL_server_ip.index(server)
                    this_node=tree.find(self.FL_server_ip[index])
                    if this_node==-1:
                        raise RuntimeError("tree_grow:{} did not found".format(server))
                    Node=tree.add_node(self.FL_server_ip[index_new],this_node)
                    tree_grow(tree,index_new,first_time)
            return
        #之前应该确保拓扑结构的完备性，即子节点与父节点的设置对应上
        if True in self.FL_server_isroot:
            root_index=self.FL_server_isroot.index(True)
        else:
            raise RuntimeError("root node not found")
        tree_grow(self.tree,root_index,first_time)
        self.tree.start_show_topology(self.window_index)
        return
    def debug_random_traffic(self):
        traffic=np.random.rand(self.FL_server_n,self.FL_server_n)
        for i in range(0,self.FL_server_n):
            traffic[i][i]=2.0
        return traffic
    def topology_strategy(self,traffic):
        connect_list=[]#节点与延迟最低的一个节点连接
        for i in range(0,self.FL_server_n):
            connect_list.append(np.argmin(traffic[i]))
        #判断是否连通,再改
        root=np.random.randint(0,self.FL_server_n)#随机选取一个节点作为root
        up_server=[]
        down_server=[]
        for i in range(0,self.FL_server_n):
            up_server.append("")
            down_server.append([])
        def find_connection(connect_list,server):
            connect=[]
            connect.append(connect_list[server])
            for i in range(0,self.FL_server_n):
                if connect_list[i]==server:
                    connect.append(i)
            return connect
            pass
        def tree_grow(index,connect_list):#广度优先生成
            visited=set()
            q = queue.Queue()
            q.put(index) #把起始点放入队列
            visited.add(index)
            while not q.empty():
                u = q.get()
                connected=find_connection(connect_list,u)
                for v in connected:
                    if v not in visited:
                        up_server[v]=self.FL_server_ip[u]
                        down_server[u].append(self.FL_server_ip[v])
                        visited.add(v)
                        q.put(v)
            return
        tree_grow(root,connect_list)
        return up_server,down_server
    def topology_test(self):
        traffic=c.debug_random_traffic()
        up_server,down_server=c.topology_strategy(traffic)
        print(traffic)
        for i in range(0,self.FL_server_n):
            self.change_topology(self.FL_server_ip[i],up_server[i],down_server[i])

        def topology_changed_completed():
            if False in self.FL_server_topology_changed:
                return False
            return True
        while not topology_changed_completed():
            time.sleep(1)
            print("waiting...")
        time.sleep(1)#这里不等会报错,因为要等待FL节点上传拓扑结构
        self.refresh()
        return
    def receive_state(self,json_data):
        if json_data['ip'] in self.FL_server_ip:
            ip_index=self.FL_server_ip.index(json_data['ip'])
            self.FL_server_state[ip_index]=json_data['state']
            self.FL_server_isroot[ip_index]=json_data['isroot']
            self.FL_server_up[ip_index]=json_data['up_server']
            self.FL_server_down[ip_index]=json_data['down_server']
            self.FL_server_timmer[ip_index]=time.time()
            if self.tree !=None:
                node=self.tree.find(json_data['ip'])
                if node!=-1:
                    node.state=json_data['state']
                    self.tree.refresh()#刷新状态显示
        else:#新加入的节点
            self.FL_server_n+=1
            self.FL_server_ip.append(json_data['ip'])
            self.FL_server_state.append(json_data['state'])
            self.FL_server_isroot.append(json_data['isroot'])
            self.FL_server_up.append(json_data['up_server'])
            self.FL_server_down.append(json_data['down_server'])
            self.FL_server_timmer.append(time.time())
            self.FL_server_topology_changed.append(False)
            if self.flag_initialized:#运行过程中新加入的节点
                print("Add new node:{}".format(json_data['ip']))
                while not self.check_state():
                    print("\rwaite...",end='')
                    time.sleep(1)
                self.topology_test()
            elif(self.check_topology()):#第一次收到完整的拓扑结构
                self.flag_initialized=True
                print("complete topology, start to show")
                self.refresh()
        return
    def check_topology(self):#结构正常返回真
        assert self.FL_server_n==len(self.FL_server_ip) , "number of servers is wrong"
        root_number=0
        for one in self.FL_server_isroot:
            if one:
                root_number+=1
        if root_number!=1:
            return False   #根节点数量不对

        for i in range(self.FL_server_n):
            up_server=self.FL_server_up[i]
            if up_server!="":
                if up_server not in self.FL_server_ip:
                    return False #缺失父节点
                up_index=self.FL_server_ip.index(up_server)
                if self.FL_server_ip[i] not in self.FL_server_down[up_index]:
                    return False #父节点没有记录该子节点

            down_servers=self.FL_server_down[i]
            for j in range(len(self.FL_server_down[i])):
                down_server=self.FL_server_down[i][j]
                if down_server not in self.FL_server_ip:
                    return False#缺失子节点
                down_index=self.FL_server_ip.index(down_server)
                if self.FL_server_up[down_index]!=self.FL_server_ip[i]:
                    return False#子节点的父节点不是该节点
        return True
    def check_state(self):
        for one in self.FL_server_state:
            if one !=1:
                return False
        return True

    def multipule_test(self):
        run_turns=20
        root_index=self.FL_server_isroot.index(True)
        self.command_send((self.FL_server_ip[root_index],self.port),self.sendjson1)

        def check_complete():#检测是否所有节点进入状态1
            for one in self.FL_server_state:
                if one != 1:
                    return False
            return True

        while(run_turns>0):
            if check_complete():
                time.sleep(3)
                if check_complete():
                    print("still have {} rounds".format(run_turns))
                    root_index=self.FL_server_isroot.index(True)
                    run_turns-=1
                    self.command_send((self.FL_server_ip[root_index],self.port),self.sendjson3)#检测模型
                    time.sleep(10)
                    self.command_send((self.FL_server_ip[root_index],self.port),self.sendjson1)#开始训练
            time.sleep(5)
        return
    #--------------基本工具-----------
    def udp_send(self,address,message):
        client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        client_socket.sendto(message, address) #将msg内容发送给指定接收方
        client_socket.close();
        return
    def command_send(self,address,message):#需要回执的命令 成功则返回true
        message_dic=json.loads(message)
        random_id=np.random.randint(1e6)
        if "feedback_id" in message_dic or "addr" in message_dic:
            raise RuntimeWarning("feedback_id addr in message_dic is overrode")
        message_dic['feedback_id']=random_id
        message_dic['addr']=(self.local_ip,self.port)

        message=json.dumps(message_dic)
        self.udp_send(address,message.encode("utf-8"))
        start_time=time.time()
        check_time=0.1#检查是否收到回执的间隙
        resend_time=0.5#重发消息的间隙
        resend_n=round(resend_time/check_time)
        n=0
        while(time.time()-start_time<3):
            time.sleep(check_time)
            n+=1
            if random_id in self.command_id:
                return True
            elif n==resend_n:
                self.udp_send(address,message.encode("utf-8"))#没有相应，重发命令
                n=0
        print("node {} is unresponsive, task id:{}, message{}".format(address,random_id,message))
        return False
    def feedback(self,json_data): #已经处理过的命令返回False，新命令返回True
        id=json_data["feedback_id"]
        ip=json_data["addr"][0]
        port=json_data["addr"][1]
        if "order" in json_data:#是feedback发送来的，不是新命令
            if json_data["order"]=="feedback":
                self.command_id.add(id)
                return False#回执不需要处理
        #可能是新命令
        send_dic={"order":"feedback","feedback_id":id}
        send_dic['addr']=(self.local_ip,self.port)
        sned_json=json.dumps(send_dic)
        self.udp_send((ip,port),sned_json.encode('utf-8'))
        if id in self.command_id:
            return False
        self.command_id.add(id)
        return True


c=FLcloud()
c.start()
